//
//  ImageStitcher.swift
//  JKTool
//
//  Created by 姜奎 on 2023/6/6.
//

import Foundation
import CoreGraphics
import ImageIO

extension JKTool {
    struct ImageStitcher: ParsableCommand {
        static var configuration = CommandConfiguration(
            commandName: "imageStitcher",
            _superCommandName: "jktool",
            abstract: "图片拼接",
            version: "1.0.0",
            subcommands: [Vertical.self,Horizontal.self,Square.self]
        )
            
    }
}


extension JKTool.ImageStitcher {
    struct Vertical: ParsableCommand {
        static var configuration = CommandConfiguration(
            commandName: "vertical",
            _superCommandName: "imageStitcher",
            abstract: "图片纵向拼接",
            version: "1.0.0"
        )
        
        @Option(name: .shortAndLong, help: "输出图片路径")
        var toPath: String
        
        @Option(name: .shortAndLong, help: "图片数组")
        var paths: [String]

        mutating func run() {
            
            
            let images = paths.compactMap { path in
                let imageURL = URL(fileURLWithPath: path)
                
                guard let imageSource = CGImageSourceCreateWithURL(imageURL as CFURL, nil),
                let image = CGImageSourceCreateImageAtIndex(imageSource, 0, nil) else {
                    fatalError("Could not create image source")
                }
                return image
            }
            
            var maxWidth = 0
            var maxHeight = 0
            var widthSum = 0
            var heightSum = 0
            
            for image in images {
                let width = image.width
                let height = image.height
                    
                widthSum += width
                heightSum += height
                maxWidth = max(maxWidth, width)
                maxHeight = max(maxHeight, height)
            }
            
            func v(images:[CGImage]){
                var x = 0
                var y = 0
                guard let context = CGContext(data: nil, width: maxWidth, height: heightSum, bitsPerComponent: 8, bytesPerRow: 0, space: CGColorSpaceCreateDeviceRGB(), bitmapInfo: CGImageAlphaInfo.premultipliedLast.rawValue) else {
                    return po(tip: "图片处理异常",type: .error)
                }
                
                for image in images {
                    context.draw(image, in: CGRect(x: x, y: y, width: image.width, height: image.height))
                    y += image.height
                }
                
                // 保存图像到磁盘
                let outputURL = URL(fileURLWithPath: "\(toPath)/vertical.png")
                guard let outputImage = context.makeImage() else {
                    return po(tip: "图片处理异常",type: .error)
                }
                let destination = CGImageDestinationCreateWithURL(outputURL as CFURL, kUTTypePNG, 1, nil)!
                CGImageDestinationAddImage(destination, outputImage, nil)
                CGImageDestinationFinalize(destination)
            }
            
            
            v(images: images)
            
        }
            
    }
}

extension JKTool.ImageStitcher {
    struct Horizontal: ParsableCommand {
        static var configuration = CommandConfiguration(
            commandName: "horizontal",
            _superCommandName: "imageStitcher",
            abstract: "图片横向拼接",
            version: "1.0.0"
        )
        
        @Option(name: .shortAndLong, help: "输出图片路径")
        var toPath: String
        
        @Option(name: .shortAndLong, help: "图片数组")
        var paths: [String]

        mutating func run() {
            
            
            let images = paths.compactMap { path in
                let imageURL = URL(fileURLWithPath: path)
                
                guard let imageSource = CGImageSourceCreateWithURL(imageURL as CFURL, nil),
                let image = CGImageSourceCreateImageAtIndex(imageSource, 0, nil) else {
                    fatalError("Could not create image source")
                }
                return image
            }
            
            var maxWidth = 0
            var maxHeight = 0
            var widthSum = 0
            var heightSum = 0
            
            for image in images {
                let width = image.width
                let height = image.height
                    
                widthSum += width
                heightSum += height
                maxWidth = max(maxWidth, width)
                maxHeight = max(maxHeight, height)
            }
            
            func h(images:[CGImage]){
                var x = 0
                var y = 0
                guard let context = CGContext(data: nil, width: widthSum, height: maxHeight, bitsPerComponent: 8, bytesPerRow: 0, space: CGColorSpaceCreateDeviceRGB(), bitmapInfo: CGImageAlphaInfo.premultipliedLast.rawValue) else {
                    return po(tip: "图片处理异常",type: .error)
                }
                
                for image in images {
                    context.draw(image, in: CGRect(x: x, y: y, width: image.width, height: image.height))
                    x += image.height
                }
                
                // 保存图像到磁盘
                let outputURL = URL(fileURLWithPath: "\(toPath)/horizontal.png")
                guard let outputImage = context.makeImage() else {
                    return po(tip: "图片处理异常",type: .error)
                }
                let destination = CGImageDestinationCreateWithURL(outputURL as CFURL, kUTTypePNG, 1, nil)!
                CGImageDestinationAddImage(destination, outputImage, nil)
                CGImageDestinationFinalize(destination)
            }
            
            h(images: images)
            
        }
            
    }
}

extension JKTool.ImageStitcher {
    struct Square: ParsableCommand {
        static var configuration = CommandConfiguration(
            commandName: "square",
            _superCommandName: "imageStitcher",
            abstract: "图片方形拼接（顺序比变）",
            version: "1.0.0"
        )
        
        @Option(name: .shortAndLong, help: "输出图片路径")
        var toPath: String
        
        @Option(name: .shortAndLong, help: "图片数组")
        var paths: [String]

        mutating func run() {
            
            
            let images = paths.compactMap { path in
                let imageURL = URL(fileURLWithPath: path)
                
                guard let imageSource = CGImageSourceCreateWithURL(imageURL as CFURL, nil),
                let image = CGImageSourceCreateImageAtIndex(imageSource, 0, nil) else {
                    fatalError("Could not create image source")
                }
                return image
            }
            func mergeImages(images: [CGImage]) -> CGImage? {
                let imageCount = images.count
                let maxPerRow = Int(sqrt(Double(imageCount)))
                let rowCount = Int(ceil(Double(imageCount) / Double(maxPerRow)))
                
                func splitArrayIntoGroups(array: [CGImage], groupSize: Int) -> (totalWidth:Int,totalHeight:Int,maxHeights:[Int]) {
                    var totalWidth: Int = 0
                    var totalHeight: Int = 0
                    var maxHeights: [Int] = []
                    var currentGroup: [CGImage] = []
                    
                    for element in array {
                        currentGroup.append(element)
                        
                        if currentGroup.count == groupSize {
                            let maxHeight = currentGroup.reduce(0) { max($0, $1.height) }
                            totalWidth  = max(totalWidth, currentGroup.reduce(0) { $0 + $1.width })
                            totalHeight += maxHeight
                            maxHeights.append(maxHeight)
                            currentGroup = []
                        }
                    }
                    
                    if !currentGroup.isEmpty {
                        let maxHeight = currentGroup.reduce(0) { max($0, $1.height) }
                        totalWidth  = max(totalWidth, currentGroup.reduce(0) { $0 + $1.width })
                        totalHeight += maxHeight
                        maxHeights.append(maxHeight)
                    }
                    
                    return (totalWidth, totalHeight,maxHeights)
                }
                let (totalWidth,totalHeight,maxHeights) = splitArrayIntoGroups(array: images, groupSize: maxPerRow)

                let context = CGContext(data: nil, width: totalWidth, height: totalHeight, bitsPerComponent: 8, bytesPerRow: 0, space: CGColorSpaceCreateDeviceRGB(), bitmapInfo: CGImageAlphaInfo.premultipliedLast.rawValue)

                var x = 0
                var count = 0
                var line = Int(count/maxPerRow)
                var y = totalHeight - maxHeights[line]
                for image in images {

                    let width = image.width
                    let height = image.height
                    let rect = CGRect(x: x, y: y, width: width, height: height)
                    context?.draw(image, in: rect)

                    count += 1
                    if count  == images.count {break}
                    if count % maxPerRow == 0 {
                        x = 0
                        var line = Int(count/maxPerRow)
                        y -= maxHeights[line]
                    } else {
                        x += width
                    }
                }

                let mergedImage = context?.makeImage()
                return mergedImage
            }
            


            // 保存图像到磁盘
            let outputURL = URL(fileURLWithPath: "\(toPath)/square.png")
            guard let outputImage = mergeImages(images: images) else {
                return po(tip: "图片处理异常",type: .error)
            }
            let destination = CGImageDestinationCreateWithURL(outputURL as CFURL, kUTTypePNG, 1, nil)!
            CGImageDestinationAddImage(destination, outputImage, nil)
            CGImageDestinationFinalize(destination)
            
        }
            
    }
}
